import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from nlpx.model import TextCNN
from nlpx.model.classifier import TextCNNClassifier
from nlpx.model.wrapper import FastClassifyModelWrapper, TextModelWrapper, PaddingTextModelWrapper, SplitTextModelWrapper, \
	SplitPaddingTextModelWrapper
from nlpx.tokenize import PaddingTokenizer
from nlpx.tokenize.utils import get_df_text_labels

pretrained_path = r'/Users/summy/data/albert_small_zh'


if __name__ == '__main__':
	df = pd.read_csv('~/project/python/parttime/归档/text_gcn/data/北方地区不安全事件统计20240331.csv', encoding='GBK')
	texts, labels, classes = get_df_text_labels(df, text_col='故障描述', label_col='故障标志')
	print(np.bincount(labels))
	print(1.0 / np.bincount(labels))
	

	# tokenize_vec = AlbertTokenizeVec(pretrained_path)
	# model = TextCNN(tokenize_vec.hidden_size, out_features=len(classes))
	
	tokenize_vec = PaddingTokenizer(texts=texts, cut_type='char')
	model = TextCNNClassifier(100, vocab_size=tokenize_vec.vocab_size, num_classes=len(classes))
	
	# # SplitTextModelWrapper准确率能达到1，原因待查
	# model_wrapper = SplitTextModelWrapper(model, tokenize_vec, classes=classes)
	# model_wrapper.train_evaluate(df['故障描述'].values, labels, n_jobs=3, early_stopping_rounds=1, val_size=0.2)
	
	model_wrapper = SplitPaddingTextModelWrapper(model, tokenize_vec, classes=classes)
	model_wrapper.train_evaluate(df['故障描述'].values, labels, sampler=None, early_stopping_rounds=2, amp=True)

	tests = ['昆明航B737飞机执行昆明至西安航班 西安机场进近过程中  机组将飞机襟翼位置设置错误触发近地警告事件',
			 '海航 B737飞机执行汉中至广州航班 飞机起飞后发现汉中过站期间货舱有730KG货物未卸机']

	# result = model_wrapper.predict_classes_proba(tests)
	# print(result)

	# print(model_wrapper.acc(tests, [4, 5]))

	# X_train, X_val, y_train, y_val = train_test_split(texts, labels, test_size=0.2, random_state=42)
	# model_wrapper = PaddingTextModelWrapper(model, tokenize_vec, classes=classes)
	# sampler = model_wrapper.get_balance_sampler(y_train)
	# model_wrapper.train_evaluate(X_train, y_train, (X_val, y_val), sampler=sampler, early_stopping_rounds=2)

	# model_wrapper.evaluate(texts[:100], labels[:100])
	# model_wrapper.classification_report(texts[:100], labels[:100])

	# print(model_wrapper.acc(texts[:100], labels[:100]))
